'''
SESSION (id: session_id)
|
|-- TRACE
|-- TRACE
|-- TRACE
|-- ...
'''
from datetime import datetime

from dotenv import load_dotenv, find_dotenv
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langfuse.decorators import observe, langfuse_context

_ = load_dotenv(find_dotenv())
now = datetime.now()
llm = ChatOpenAI()
messages = [
    SystemMessage(content="你是AGIClass的课程助理。"),
]


@observe()
def chat_one_turn(user_input, user_id, turn_id):
    langfuse_context.update_current_trace(
        name=f"ChatTurn{turn_id}",
        user_id=user_id,
        session_id="chat-" + now.strftime("%d/%m/%Y %H:%M:%S")
    )
    langfuse_handler = langfuse_context.get_current_langchain_handler()
    messages.append(HumanMessage(content=user_input))
    response = llm.invoke(messages, config={"callbacks": [langfuse_handler]})
    messages.append(response)
    return response.content


user_id = "wzr"
turn_id = 0
while True:
    user_input = input("User: ")
    if user_input.strip() == "":
        break
    reply = chat_one_turn(user_input, user_id, turn_id)
    print("AI: " + reply)
    turn_id += 1
